
#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
using namespace std;
using namespace Eigen;

struct Fun {
  double operator()(const VectorXd &x) const {
    // return 5.0 + x(0) * x(0) * x(1) * sin(x(1) / 10) / 10000;
    return (1.0 + 10 * x(0) + 3 * x(1) + x(0) * x(0) + x(1) * x(1) -
            2 * x(0) * x(1)) /
           1000;
  }
};

// 接收一个采样点矩阵,计算出系数
template <typename T> struct QuadFun {
  VectorXd coe;
  int nv;
  QuadFun(){};

  void calcCoe(T f, const MatrixXd &sam_point) {
    int row = nv = sam_point.rows();
    int col = sam_point.cols();
    coe.resize((row + 1) * (row + 2) / 2);
    assert(col >= coe.size());
    MatrixXd mat1(col, coe.size());
    VectorXd y(col);
    // 计算系数矩阵和函数值
    VectorXd tmp(row);
    for (int i = 0; i < col; ++i) {
      mat1(i, 0) = 1.0;
      int count = 1;
      for (int j = 0; j < row; ++j) {
        mat1(i, j + 1) = sam_point(j, i);
        for (int k = j; k < row; ++k) {
          mat1(i, row + count) = sam_point(j, i) * sam_point(k, i);
          ++count;
        }
      }
      assert(count == 1 + (row * row + row) / 2);
      // 计算函数值
      //   tmp = sam_point.col(i);
      for (int h = 0; h < row; ++h)
        tmp(h) = sam_point(h, i);
      y(i) = f(tmp);
    }

    MatrixXd mat2 = mat1.transpose() * mat1;
    VectorXd y2 = mat1.transpose() * y;

    coe = mat2.colPivHouseholderQr().solve(y2);
  }

  double operator()(const VectorXd &x) {
    assert(nv == x.size());
    double y = coe(0);

    int count = 1;
    for (int j = 0; j < nv; ++j) {
      y += coe(j + 1) * x(j);
      for (int k = j; k < nv; ++k) {
        // mat1(i, row + count) = sam_point(j, i) * sam_point(k, i);
        y += coe(nv + count) * x(j) * x(k);
        ++count;
      }
    }
    return y;
  }
};

int main() {
  const int range = 100;
  Fun f;
  ofstream of1{"real_fun.dat"};
  of1 << scientific << setprecision(15);
  of1 << " variables = \"x\", \"y\", \"f\"\n";
  of1 << "ZONE I = 100, J = 100\n";
  VectorXd tmp(2);
  for (int i = 1; i <= range; ++i)
    for (int j = 1; j <= range; ++j) {
      tmp(0) = i;
      tmp(1) = j;
      of1 << tmp(0) << "  " << tmp(1) << "  " << f(tmp) << "\n";
    }
  of1.close();

  // 全因子设计
  const int gap{10}, np{range / gap};
  MatrixXd sam_point(2, np * np);
  int count{0};
  for (int i = 1; i <= range; i += gap) {
    if ((i - 1) / gap == np)
      break;
    for (int j = 1; j <= range; j += gap) {
      if ((i - 1) / gap == np)
        break;
      sam_point(0, count) = i;
      sam_point(1, count) = j;
      count++;
    }
  }
  cout << sam_point << endl;
  assert(count == np * np);
  QuadFun<Fun> qf;
  qf.calcCoe(f, sam_point);

  cout << setprecision(15);
  for (int i = 0; i < qf.coe.size(); ++i)
    cout << qf.coe(i) << "   ";
  cout << endl;

  ofstream of2{"qf_fun.dat"};
  of2 << scientific << setprecision(15);
  of2 << " variables = \"x\", \"y\", \"qf\"\n";
  of2 << "ZONE I = 100, J = 100\n";
  //   VectorXd tmp2(2);
  for (int i = 1; i <= range; ++i)
    for (int j = 1; j <= range; ++j) {
      tmp(0) = i;
      tmp(1) = j;
      of2 << tmp(0) << "  " << tmp(1) << "  " << qf(tmp) << "\n";
    }
  of2.close();

  return 0;
}